#include <my_dwa_planner/dwa_planner.h>
#include <pluginlib/class_list_macros.h>
#include <cmath>
#include <algorithm>
#include <Eigen/Core>

PLUGINLIB_EXPORT_CLASS(my_dwa_planner::DWAPlanner, nav_core::BaseLocalPlanner)

namespace my_dwa_planner {

DWAPlanner::DWAPlanner() : 
    costmap_ros_(nullptr), tf_(nullptr), world_model_(nullptr), dsrv_(nullptr), 
    goal_reached_(false), initialized_(false) {
}

DWAPlanner::~DWAPlanner() {
    if(dsrv_) delete dsrv_;
    if(world_model_) delete world_model_;
}

void DWAPlanner::initialize(std::string name, tf2_ros::Buffer* tf, costmap_2d::Costmap2DROS* costmap_ros) {
    if(!initialized_) {
        tf_ = tf;
        costmap_ros_ = costmap_ros;
        
        ros::NodeHandle private_nh("~/" + name);
        ros::NodeHandle nh;
        
        dsrv_ = new dynamic_reconfigure::Server<DWAPlannerConfig>(private_nh);
        dynamic_reconfigure::Server<DWAPlannerConfig>::CallbackType cb = 
            boost::bind(&DWAPlanner::reconfigureCB, this, _1, _2);
        dsrv_->setCallback(cb);
        
        world_model_ = new base_local_planner::CostmapModel(*costmap_ros_->getCostmap());
        
        traj_pub_ = private_nh.advertise<nav_msgs::Path>("trajectories", 1);
        best_traj_pub_ = private_nh.advertise<nav_msgs::Path>("best_trajectory", 1);
        
        initialized_ = true;
        ROS_INFO("DWA planner initialized");
    }
}

bool DWAPlanner::setPlan(const std::vector<geometry_msgs::PoseStamped>& plan) {
    if(!initialized_) {
        ROS_ERROR("Planner has not been initialized");
        return false;
    }
    
    global_plan_ = plan;
    goal_reached_ = false;
    return true;
}

bool DWAPlanner::computeVelocityCommands(geometry_msgs::Twist& cmd_vel) {
    if(!initialized_) {
        ROS_ERROR("Planner has not been initialized");
        return false;
    }
    
    geometry_msgs::PoseStamped robot_pose;
    if(!costmap_ros_->getRobotPose(robot_pose)) {
        ROS_ERROR("Could not get robot pose");
        return false;
    }
    
    nav_msgs::Odometry robot_vel;
    // 实际实现中应从odom话题获取
    
    std::vector<geometry_msgs::PoseStamped> transformed_plan;
    if(!transformGlobalPlan(*tf_, global_plan_, *costmap_ros_, costmap_ros_->getGlobalFrameID(), transformed_plan)) {
        ROS_WARN("Could not transform the global plan");
        return false;
    }
    
    if(transformed_plan.empty()) {
        ROS_WARN("Transformed plan is empty");
        return false;
    }
    
    const geometry_msgs::PoseStamped& goal_pose = transformed_plan.back();
    double dx = goal_pose.pose.position.x - robot_pose.pose.position.x;
    double dy = goal_pose.pose.position.y - robot_pose.pose.position.y;
    double distance = hypot(dx, dy);
    
    if(distance <= dwa_xy_tolerance_) {
        goal_reached_ = true;
        cmd_vel.linear.x = 0.0;
        cmd_vel.linear.y = 0.0;
        cmd_vel.angular.z = 0.0;
        return true;
    }
    
    geometry_msgs::PoseStamped best_pose;
    base_local_planner::Trajectory best_traj = findBestPath(robot_pose, robot_vel, best_pose);
    
    if(best_traj.cost_ < 0) {
        ROS_WARN("Could not find valid trajectory");
        return false;
    }
    
    nav_msgs::Path best_path;
    best_path.header.stamp = ros::Time::now();
    best_path.header.frame_id = costmap_ros_->getGlobalFrameID();
    geometry_msgs::PoseStamped pose;
    pose.header = best_path.header;
    
    for(unsigned int i = 0; i < best_traj.getPointsSize(); ++i) {
        double p_x, p_y, p_th;
        best_traj.getPoint(i, p_x, p_y, p_th);
        pose.pose.position.x = p_x;
        pose.pose.position.y = p_y;
        tf2::Quaternion q;
        q.setRPY(0, 0, p_th);
        tf2::convert(q, pose.pose.orientation);
        best_path.poses.push_back(pose);
    }
    best_traj_pub_.publish(best_path);
    
    cmd_vel.linear.x = best_traj.xv_;
    cmd_vel.linear.y = best_traj.yv_;
    cmd_vel.angular.z = best_traj.thetav_;
    
    return true;
}

bool DWAPlanner::isGoalReached() {
    if(!initialized_) {
        ROS_ERROR("Planner has not been initialized");
        return false;
    }
    return goal_reached_;
}

base_local_planner::Trajectory DWAPlanner::findBestPath(
    const geometry_msgs::PoseStamped& robot_pose,
    const nav_msgs::Odometry& robot_vel,
    geometry_msgs::PoseStamped& best_pose) {
    
    double vx_min = std::max(min_vel_x_, robot_vel.twist.twist.linear.x - acc_lim_x_ * sim_time_);
    double vx_max = std::min(max_vel_x_, robot_vel.twist.twist.linear.x + acc_lim_x_ * sim_time_);
    
    double vy_min = std::max(min_vel_y_, robot_vel.twist.twist.linear.y - acc_lim_y_ * sim_time_);
    double vy_max = std::min(max_vel_y_, robot_vel.twist.twist.linear.y + acc_lim_y_ * sim_time_);
    
    double vth_min = std::max(min_vel_th_, robot_vel.twist.twist.angular.z - acc_lim_th_ * sim_time_);
    double vth_max = std::min(max_vel_th_, robot_vel.twist.twist.angular.z + acc_lim_th_ * sim_time_);
    
    std::vector<base_local_planner::Trajectory> all_explored;
    base_local_planner::Trajectory best_traj;
    best_traj.cost_ = -1;
    
    double vx_samp = vx_min;
    while(vx_samp <= vx_max) {
        double vy_samp = vy_min;
        while(vy_samp <= vy_max) {
            double vth_samp = vth_min;
            while(vth_samp <= vth_max) {
                base_local_planner::Trajectory traj;
                double x_i = robot_pose.pose.position.x;
                double y_i = robot_pose.pose.position.y;
                double theta_i = tf2::getYaw(robot_pose.pose.orientation);
                
                traj.xv_ = vx_samp;
                traj.yv_ = vy_samp;
                traj.thetav_ = vth_samp;
                traj.cost_ = 0.0;
                
                double time = 0.0;
                while(time < sim_time_) {
                    x_i += vx_samp * cos(theta_i) * sim_granularity_;
                    y_i += vx_samp * sin(theta_i) * sim_granularity_;
                    theta_i += vth_samp * sim_granularity_;
                    
                    traj.addPoint(x_i, y_i, theta_i);
                    time += sim_granularity_;
                }
                
                traj.cost_ = scoreTrajectory(traj);
                all_explored.push_back(traj);
                
                if(traj.cost_ >= 0 && (traj.cost_ < best_traj.cost_ || best_traj.cost_ < 0)) {
                    best_traj = traj;
                    best_pose.pose.position.x = x_i;
                    best_pose.pose.position.y = y_i;
                    tf2::Quaternion q;
                    q.setRPY(0, 0, theta_i);
                    tf2::convert(q, best_pose.pose.orientation);
                }
                
                vth_samp += angular_sim_granularity_;
            }
            vy_samp += sim_granularity_;
        }
        vx_samp += sim_granularity_;
    }
    
    nav_msgs::Path traj_path;
    traj_path.header.stamp = ros::Time::now();
    traj_path.header.frame_id = costmap_ros_->getGlobalFrameID();
    
    for(unsigned int i = 0; i < all_explored.size(); ++i) {
        geometry_msgs::PoseStamped pose;
        pose.header = traj_path.header;
        double x, y, th;
        all_explored[i].getEndpoint(x, y, th);
        pose.pose.position.x = x;
        pose.pose.position.y = y;
        tf2::Quaternion q;
        q.setRPY(0, 0, th);
        tf2::convert(q, pose.pose.orientation);
        traj_path.poses.push_back(pose);
    }
    traj_pub_.publish(traj_path);
    
    return best_traj;
}

double DWAPlanner::scoreTrajectory(const base_local_planner::Trajectory& traj) {
    double cost = 0.0;
    double footprint_cost = 0.0;
    
    for(unsigned int i = 0; i < traj.getPointsSize(); ++i) {
        double x, y, th;
        traj.getPoint(i, x, y, th);
        footprint_cost = world_model_->footprintCost(x, y, th, costmap_ros_->getRobotFootprint(), 0.0, 0.0);
        
        if(footprint_cost < 0) {
            return -1.0;
        }
        cost += footprint_cost * occdist_scale_;
    }
    
    double path_dist = 0.0;
    double goal_dist = 0.0;
    double heading_diff = 0.0;
    
    double x, y, th;
    traj.getEndpoint(x, y, th);
    
    double min_dist = std::numeric_limits<double>::max();
    int min_index = 0;
    for(unsigned int i = 0; i < global_plan_.size(); ++i) {
        double dx = x - global_plan_[i].pose.position.x;
        double dy = y - global_plan_[i].pose.position.y;
        double dist = hypot(dx, dy);
        if(dist < min_dist) {
            min_dist = dist;
            min_index = i;
        }
    }
    
    path_dist = min_dist;
    
    double dx = x - global_plan_.back().pose.position.x;
    double dy = y - global_plan_.back().pose.position.y;
    goal_dist = hypot(dx, dy);
    
    if(min_index + heading_lookahead_ < global_plan_.size()) {
        heading_diff = headingDiff(x, y, 
                                 global_plan_[min_index + heading_lookahead_].pose.position.x,
                                 global_plan_[min_index + heading_lookahead_].pose.position.y);
    } else {
        heading_diff = headingDiff(x, y, 
                                 global_plan_.back().pose.position.x,
                                 global_plan_.back().pose.position.y);
    }
    
    cost += path_dist * pdist_scale_;
    cost += goal_dist * gdist_scale_;
    cost += fabs(heading_diff) * 0.5;
    
    return cost;
}

double DWAPlanner::headingDiff(double x, double y, double pt_x, double pt_y) {
    double v1_x = pt_x - x;
    double v1_y = pt_y - y;
    double v2_x = cos(0); // 机器人当前朝向
    double v2_y = sin(0);
    
    double perp_dot = v1_x * v2_y - v1_y * v2_x;
    double dot = v1_x * v2_x + v1_y * v2_y;
    
    return atan2(perp_dot, dot);
}

void DWAPlanner::reconfigureCB(DWAPlannerConfig &config, uint32_t level) {
    max_vel_x_ = config.max_vel_x;
    min_vel_x_ = config.min_vel_x;
    max_vel_y_ = config.max_vel_y;
    min_vel_y_ = config.min_vel_y;
    max_vel_th_ = config.max_vel_th;
    min_vel_th_ = config.min_vel_th;
    
    acc_lim_x_ = config.acc_lim_x;
    acc_lim_y_ = config.acc_lim_y;
    acc_lim_th_ = config.acc_lim_th;
    
    sim_time_ = config.sim_time;
    sim_granularity_ = config.sim_granularity;
    angular_sim_granularity_ = config.angular_sim_granularity;
    
    pdist_scale_ = config.pdist_scale;
    gdist_scale_ = config.gdist_scale;
    occdist_scale_ = config.occdist_scale;
    
    stop_time_buffer_ = config.stop_time_buffer;
    oscillation_reset_dist_ = config.oscillation_reset_dist;
    heading_lookahead_ = config.heading_lookahead;
    
    dwa_xy_tolerance_ = config.xy_goal_tolerance;
    dwa_yaw_tolerance_ = config.yaw_goal_tolerance;
}

bool DWAPlanner::transformGlobalPlan(const tf2_ros::Buffer& tf, 
                                   const std::vector<geometry_msgs::PoseStamped>& global_plan,
                                   const costmap_2d::Costmap2DROS& costmap,
                                   const std::string& global_frame,
                                   std::vector<geometry_msgs::PoseStamped>& transformed_plan) {
    const geometry_msgs::PoseStamped& plan_pose = global_plan[0];
    transformed_plan.clear();
    
    try {
        if(global_plan.empty()) {
            ROS_ERROR("Received plan with zero length");
            return false;
        }
        
        geometry_msgs::TransformStamped transform = tf.lookupTransform(
            global_frame, ros::Time(), 
            plan_pose.header.frame_id, plan_pose.header.stamp,
            plan_pose.header.frame_id);
        
        geometry_msgs::PoseStamped newer_pose;
        unsigned int i = 0;
        double sq_dist_threshold = 1.0;
        double sq_dist = 0.0;
        
        for(i = 0; i < global_plan.size(); ++i) {
            const geometry_msgs::PoseStamped& pose = global_plan[i];
            tf2::doTransform(pose, newer_pose, transform);
            
            double x = newer_pose.pose.position.x;
            double y = newer_pose.pose.position.y;
            
            unsigned int cell_x, cell_y;
            if(!costmap.getCostmap()->worldToMap(x, y, cell_x, cell_y)) {
                break;
            }
            
            if(i > 0) {
                double dx = newer_pose.pose.position.x - transformed_plan.back().pose.position.x;
                double dy = newer_pose.pose.position.y - transformed_plan.back().pose.position.y;
                sq_dist = dx*dx + dy*dy;
                
                if(sq_dist > sq_dist_threshold) {
                    break;
                }
            }
            
            transformed_plan.push_back(newer_pose);
        }
    } catch(tf2::LookupException& ex) {
        ROS_ERROR("No Transform available: %s", ex.what());
        return false;
    } catch(tf2::ConnectivityException& ex) {
        ROS_ERROR("Connectivity Error: %s", ex.what());
        return false;
    } catch(tf2::ExtrapolationException& ex) {
        ROS_ERROR("Extrapolation Error: %s", ex.what());
        return false;
    }
    
    if(transformed_plan.empty()) {
        ROS_ERROR("Resulting plan has 0 poses in it.");
        return false;
    }
    
    return true;
}

} // namespace my_dwa_planner